import numpy as np
from numpy.random import choice
import random

def get_flat_block(cfg, heightfield_template, block_length, block_width, difficulty = 0.1,block_starting_height_px = 0):
    # goal in the middle of the block
    goals = np.array([[block_length / 2 * cfg.horizontal_scale, block_width / 2 * cfg.horizontal_scale, 0]])
    
    track_heightfield = heightfield_template.copy()  
    if block_starting_height_px != 0:
        track_heightfield += round(block_starting_height_px)
    height_offset_px = 0
    if random.random() < 0.99:
        return track_heightfield, track_heightfield.copy(), None, height_offset_px, goals   
    else:
        # max_height = np.random.uniform(*cfg.discrete["max_height"]) if isinstance(cfg.discrete["max_height"], (tuple, list)) else cfg.discrete["max_height"]
        if isinstance(cfg.discrete["max_height"], (tuple, list)):
            if difficulty is None:
                max_height = np.random.uniform(*cfg.discrete["max_height"])
            else:
                max_height = (1-difficulty) * cfg.discrete["max_height"][0] + difficulty * cfg.discrete["max_height"][1]
        else:
            max_height = cfg.discrete["max_height"]
        max_height = round(max_height/cfg.vertical_scale)
        
        min_size = np.random.uniform(*cfg.discrete["min_size"]) if isinstance(cfg.discrete["min_size"], (tuple, list)) else cfg.discrete["min_size"]
        min_size = round(min_size/cfg.horizontal_scale)
        max_size = np.random.uniform(*cfg.discrete["max_size"]) if isinstance(cfg.discrete["max_size"], (tuple, list)) else cfg.discrete["max_size"]
        max_size = round(max_size/cfg.horizontal_scale)

        (i, j) = track_heightfield.shape
        height_range = [-max_height, -max_height // 2, max_height // 2, max_height]
        width_range = range(min_size, max_size, 4)
        length_range = range(min_size, max_size, 4)
        
        for _ in range(cfg.discrete["num_rects"]):
            width = np.random.choice(width_range)
            length = np.random.choice(length_range)
            start_i = np.random.choice(range(0, i-width, 4))
            start_j = np.random.choice(range(0, j-length, 4))
            track_heightfield[start_i:start_i+width, start_j:start_j+length] += np.random.choice(height_range)
        return track_heightfield, track_heightfield.copy(), None, height_offset_px, goals   

def get_hurdle_block(cfg, heightfield_template, block_length, block_width, difficulty = 0.1, block_starting_height_px = 0):
    goals = np.zeros((2, 3))
    track_heightfield = heightfield_template.copy()  
    if block_starting_height_px != 0:
        track_heightfield += round(block_starting_height_px)
    
    hurdle_length = np.random.uniform(*cfg.hurdle["length"]) if isinstance(cfg.hurdle["length"], (tuple, list)) else cfg.hurdle["length"]
    if isinstance(cfg.hurdle["length"], (tuple, list)):
        if difficulty is None:
            hurdle_height = np.random.uniform(*cfg.hurdle["height"])
        else:
            hurdle_height = (1-difficulty) * cfg.hurdle["height"][0] + difficulty * cfg.hurdle["height"][1]
    else:
        hurdle_height = cfg.hurdle["height"]
    length_px = (hurdle_length / cfg.horizontal_scale)
    height_value = hurdle_height / cfg.vertical_scale
    
    offset = np.random.uniform(*cfg.hurdle["fake_offset"]) if \
            isinstance(cfg.hurdle["fake_offset"], (tuple, list)) else \
            np.random.uniform(-cfg.hurdle["fake_offset"], cfg.hurdle["fake_offset"],)
    start_px = int(offset / cfg.horizontal_scale)
    
    track_heightfield[
        int(start_px + track_heightfield.shape[0] / 2) - round(length_px/2): int(start_px + track_heightfield.shape[0] / 2) + round(length_px/2), :,
    ] += round(height_value)
    mid_y = track_heightfield.shape[1] / 2
    
    goals[0, 0] = (start_px + (track_heightfield.shape[0]-length_px) / 2) * cfg.horizontal_scale - cfg.hurdle["goal_offset"]
    goals[1, 0] = (start_px + (track_heightfield.shape[0]+length_px) / 2) * cfg.horizontal_scale + cfg.hurdle["goal_offset"]
    goals[:, 1] = mid_y * cfg.horizontal_scale
    # index_x = np.clip((goals[:, 0] / self.cfg.horizontal_scale).astype(np.int), a_max= heightfield_template.shape[0]-1, a_min= 0)
    # index_y = np.clip((goals[:, 1] / self.cfg.horizontal_scale).astype(np.int), a_max= heightfield_template.shape[1]-1, a_min= 0)
    # goals[:, 2] = (heightfield_template[index_x, index_y] + block_starting_height_px) * self.cfg.vertical_scale
    
    # block_info = torch.tensor([
    #     hurdle_length + self.track_kwargs["hurdle"].get("fake_offset", 0.), # along x(forward)-axis
    #     hurdle_depth, # along z(downward)-axis
    # ], dtype= torch.float32, device= self.device)
    height_offset_px = 0
    return track_heightfield, track_heightfield.copy(), None, height_offset_px, goals

def get_gap_block(cfg, heightfield_template, block_length, block_width, difficulty = 0.1,block_starting_height_px = 0):
    goals = np.zeros((2, 3))
    track_heightfield = heightfield_template.copy()
    if block_starting_height_px != 0:
        track_heightfield += round(block_starting_height_px)
        
    gap_depth = np.random.uniform(*cfg.gap["depth"]) if isinstance(cfg.gap["depth"], (tuple, list)) else cfg.gap["depth"]
    if isinstance(cfg.gap["length"], (tuple, list)):
        if difficulty is None:
            gap_length = np.random.uniform(*cfg.gap["length"])
        else:
            gap_length = (1-difficulty) * cfg.gap["length"][0] + difficulty * cfg.gap["length"][1]
    else:
        gap_length = cfg.gap["length"]
    length_px = (gap_length / cfg.horizontal_scale)
    depth_value = -gap_depth / cfg.vertical_scale

    offset = np.random.uniform(*cfg.gap["fake_offset"]) if \
            isinstance(cfg.gap["fake_offset"], (tuple, list)) else \
            np.random.uniform(-cfg.gap["fake_offset"], cfg.gap["fake_offset"],)
    start_px = int(offset / cfg.horizontal_scale)
    
    track_heightfield[
        int(start_px + track_heightfield.shape[0] / 2) - round(length_px/2): int(start_px + track_heightfield.shape[0] / 2) + round(length_px/2), :,
    ] = round(depth_value)
    mid_y = track_heightfield.shape[1] / 2
    
    goals[0, 0] = (start_px + (track_heightfield.shape[0]-length_px) / 2) * cfg.horizontal_scale - cfg.gap["goal_offset"]
    goals[1, 0] = (start_px + (track_heightfield.shape[0]+length_px) / 2) * cfg.horizontal_scale + cfg.gap["goal_offset"]
    goals[:, 1] = mid_y * cfg.horizontal_scale
    # index_x = np.clip((goals[:, 0] / self.cfg.horizontal_scale).astype(np.int), a_max= heightfield_template.shape[0]-1, a_min= 0)
    # index_y = np.clip((goals[:, 1] / self.cfg.horizontal_scale).astype(np.int), a_max= heightfield_template.shape[1]-1, a_min= 0)
    # goals[:, 2] = (heightfield_template[index_x, index_y] + block_starting_height_px) * self.cfg.vertical_scale
    
    # block_info = torch.tensor([
    #     hurdle_length + self.track_kwargs["hurdle"].get("fake_offset", 0.), # along x(forward)-axis
    #     hurdle_depth, # along z(downward)-axis
    # ], dtype= torch.float32, device= self.device)
    height_offset_px = 0
    return track_heightfield, track_heightfield.copy(), None, height_offset_px, goals

def get_jump_up_block(cfg, heightfield_template, block_length, block_width, difficulty = 0.1,block_starting_height_px = 0):
    # goal in the middle of the block
    goals = np.array([[block_length / 2 * cfg.horizontal_scale, block_width / 2 * cfg.horizontal_scale, 0]])
    
    # if isinstance(cfg.jump_up["length"], (tuple, list)):
    #     jump_length = min(*cfg.jump_up["length"])
    # else:
    #     jump_length = cfg.jump_up["length"]
        
    if isinstance(cfg.jump_up["height"], (tuple, list)):
        if difficulty is None:
            jump_height = np.random.uniform(*cfg.jump_up["height"])
        else:
            jump_height = (1-difficulty) * cfg.jump_up["height"][0] + difficulty * cfg.jump_up["height"][1]
    else:
        jump_height = cfg.jump_up["height"]
        
    height_value = jump_height / cfg.vertical_scale    
    track_heightfield = heightfield_template.copy()   
    if block_starting_height_px != 0:
        track_heightfield += round(block_starting_height_px)
    # offset = np.random.uniform(*cfg.jump_up["fake_offset"]) if \
    #         isinstance(cfg.jump_up["fake_offset"], (tuple, list)) else \
    #         np.random.uniform(-cfg.jump_up["fake_offset"], cfg.jump_up["fake_offset"],)
    # offset = int(offset / cfg.horizontal_scale)
    
    # track_heightfield[
            # int(track_heightfield.shape[0]/2) + offset:,
            # wall_thickness_px: -wall_thickness_px,
        # ] += height_value
    track_heightfield[
            : , :
        ] += round(height_value)
    # track_heightfield[
    #         :, :
    #     ] += height_value
    # mid_y = track_heightfield.shape[1] / 2
    
    height_offset_px = height_value
    return track_heightfield, track_heightfield.copy(), None, height_offset_px, goals

def get_jump_down_block(cfg, heightfield_template, block_length, block_width, difficulty = 0.1, block_starting_height_px = 0):
    # goal in the middle of the block
    goals = np.array([[block_length / 2 * cfg.horizontal_scale, block_width / 2 * cfg.horizontal_scale, 0]])
    
    # if isinstance(cfg.jump_down["length"], (tuple, list)):
    #     jump_length = min(*cfg.jump_down["length"])
    # else:
    #     jump_length = cfg.jump_down["length"]
        
    if isinstance(cfg.jump_down["depth"], (tuple, list)):
        if difficulty is None:
            jump_depth = np.random.uniform(*cfg.jump_down["depth"])
        else:
            jump_depth = (1-difficulty) * cfg.jump_down["depth"][0] + difficulty * cfg.jump_down["depth"][1]
    else:
        jump_depth = cfg.jump_down["depth"]
        
    height_value = -jump_depth / cfg.vertical_scale    
    track_heightfield = heightfield_template.copy()   
    if block_starting_height_px != 0:
        track_heightfield += round(block_starting_height_px)
    # offset = np.random.uniform(*cfg.jump_up["fake_offset"]) if \
    #         isinstance(cfg.jump_up["fake_offset"], (tuple, list)) else \
    #         np.random.uniform(-cfg.jump_up["fake_offset"], cfg.jump_up["fake_offset"],)
    # offset = int(offset / cfg.horizontal_scale)
    
    # track_heightfield[
            # int(track_heightfield.shape[0]/2) + offset:,
            # wall_thickness_px: -wall_thickness_px,
        # ] += height_value
    track_heightfield[
            : , :
        ] += round(height_value)
    # track_heightfield[
    #         :, :
    #     ] += height_value
    # mid_y = track_heightfield.shape[1] / 2
    
    height_offset_px = height_value
    return track_heightfield, track_heightfield.copy(), None, height_offset_px, goals

def get_forbidden_block(cfg, heightfield_template, block_length, block_width, difficulty = 0.1, block_starting_height_px = 0):       
    # random position   random size  random height 
    if difficulty is None:
        obs_width = np.random.uniform(*cfg.forbidden["obs_size"]) / cfg.horizontal_scale
        obs_length = np.random.uniform(*cfg.forbidden["obs_size"]) / cfg.horizontal_scale
    else:
        obs_width = (1-difficulty) * cfg.forbidden["obs_size"][0] / cfg.horizontal_scale + difficulty * cfg.forbidden["obs_size"][1] / cfg.horizontal_scale
        obs_length = (1-difficulty) * cfg.forbidden["obs_size"][0] / cfg.horizontal_scale + difficulty * cfg.forbidden["obs_size"][1] / cfg.horizontal_scale
    forbidden_height = np.random.uniform(*cfg.forbidden["height"])
    
    # sample obstacle center
    obs_center_x = np.random.uniform(obs_length+5, block_length - obs_length-6)
    obs_center_y = np.random.uniform(obs_width+5, block_width - obs_width-6)
        
    height_value = forbidden_height / cfg.vertical_scale    
    track_heightfield = heightfield_template.copy()  
    if block_starting_height_px != 0:
        track_heightfield += round(block_starting_height_px) 
    virtual_height_field = track_heightfield.copy()
    
    track_heightfield[int(obs_center_x-obs_length//2):int(obs_center_x+obs_length//2),
                    int(obs_center_y-obs_width//2):int(obs_center_y+obs_width//2),
                    ] += round(height_value)
    
    virtual_height_field[int(obs_center_x-obs_length//2)-2:int(obs_center_x+obs_length//2)-2,
                    int(obs_center_y-obs_width//2)-2:int(obs_center_y+obs_width//2)-2,
                    ] += round(height_value)

    height_offset_px = 0
    return track_heightfield, virtual_height_field, None, height_offset_px, None

def get_stairs_up_block(cfg, heightfield_template, block_length, block_width, difficulty = 0.1,block_starting_height_px = 0):
    if isinstance(cfg.stairs_up["num_steps"], (tuple, list)):
        num_steps = random.randint(*cfg.stairs_up["num_steps"])
    else:
        num_steps = cfg.stairs_up["num_steps"]
        
    if isinstance(cfg.stairs_up["step_height"], (tuple, list)):
        if difficulty is None:
            step_height = np.random.uniform(*cfg.stairs_up["step_height"])
            step_len = np.random.uniform(*cfg.stairs_up["step_len"])
        else:
            step_height = (1-difficulty) * cfg.stairs_up["step_height"][0] + difficulty * cfg.stairs_up["step_height"][1]
            step_len = (1-difficulty) * cfg.stairs_up["step_len"][0] + difficulty * cfg.stairs_up["step_len"][1]
    else:
        step_height = cfg.stairs_up["step_height"]
        step_len = cfg.stairs_up["step_len"]
        
    step_height = step_height / cfg.vertical_scale    
    step_len = step_len / cfg.horizontal_scale
    track_heightfield = heightfield_template.copy()   
    if block_starting_height_px != 0:
        track_heightfield += round(block_starting_height_px)
    
    start_x = step_len * 3; end_x = step_len * 4
    for i in range(num_steps):
        delta_height = step_height * (i)
        track_heightfield[int(start_x):int(end_x), :] += np.int16(round(delta_height))
        start_x += step_len
        end_x += step_len
        
    track_heightfield[int(start_x):, :] += np.int16(round(delta_height))
    height_offset_px = (num_steps-1) * step_height
    return track_heightfield, track_heightfield.copy(), None, height_offset_px, None

def get_stairs_down_block(cfg, heightfield_template, block_length, block_width, difficulty = 0.1,block_starting_height_px = 0):
    if isinstance(cfg.stairs_down["num_steps"], (tuple, list)):
        num_steps = random.randint(*cfg.stairs_down["num_steps"])
    else:
        num_steps = cfg.stairs_down["num_steps"]
        
    if isinstance(cfg.stairs_down["step_height"], (tuple, list)):
        if difficulty is None:
            step_height = np.random.uniform(*cfg.stairs_down["step_height"])
            step_len = np.random.uniform(*cfg.stairs_down["step_len"])
        else:
            step_height = (1-difficulty) * cfg.stairs_down["step_height"][0] + difficulty * cfg.stairs_down["step_height"][1]
            step_len = (1-difficulty) * cfg.stairs_down["step_len"][0] + difficulty * cfg.stairs_down["step_len"][1]
    else:
        step_height = cfg.stairs_down["step_height"]
        step_len = cfg.stairs_down["step_len"]
        
    step_height = -step_height / cfg.vertical_scale    
    step_len = step_len / cfg.horizontal_scale
    track_heightfield = heightfield_template.copy()   
    if block_starting_height_px != 0:
        track_heightfield += round(block_starting_height_px)
    
    start_x = step_len * 3; end_x = step_len * 4
    for i in range(num_steps):
        delta_height = step_height * (i)
        track_heightfield[int(start_x):int(end_x), :] += np.int16(round(delta_height))
        start_x += step_len
        end_x += step_len
        
    track_heightfield[int(start_x):, :] += np.int16(round(delta_height))
    height_offset_px = (num_steps-1) * step_height
    return track_heightfield, track_heightfield.copy(), None, height_offset_px, None
